import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.ensemble import IsolationForest
from matplotlib.ticker import MultipleLocator


p = "/home/duxiangyu/origindata/processed_data/"
sp = "/home/duxiangyu/origindata/labeled_data/"
for filename in os.listdir(p):
    X_full = pd.read_csv(os.path.join(p, filename))
    X_full['anomaly_label'] = None
    
    # 筛除第一类异常点，第三类异常点
    X_full.loc[X_full['Patv'] < 0, 'anomaly_label'] = -1 
    X_full.loc[(X_full['Patv'] < 1) & (X_full['Wspd'] > 2.5), 'anomaly_label'] = -1
    X_full.loc[(X_full['Patv'] < 1400) & (X_full['Wspd'] > 13), 'anomaly_label'] = -1
    # 标记好可能被误判的满发点
    X_full.loc[(X_full['Patv'] > 1400) & (X_full['Wspd'] > 8), 'anomaly_label'] = 1

    X = X_full.loc[X_full['anomaly_label'].isna(), ['Wspd', 'Patv']]

    X_low = X[X['Patv'] < 700]
    # 模型初始化
    model_low = IsolationForest(n_estimators=256, contamination=0.05, max_samples=1000, random_state=42)
    model_low.fit(X_low)
    # 预测 (-1 表示异常点，1 表示正常点)
    y_pred_low = model_low.predict(X_low)

    X_high = X[X['Patv'] >= 700]
    model_high = IsolationForest(n_estimators=256, contamination=0.05, max_samples=1000, random_state=42)
    model_high.fit(X_high)
    y_pred_high = model_high.predict(X_high)

    X_full.loc[X_low.index, 'anomaly_label'] = y_pred_low
    X_full.loc[X_high.index, 'anomaly_label'] = y_pred_high
    X_full.to_csv(os.path.join(sp, f'a_{filename}'), index=False)
    
    # 可视化
    plt.figure(figsize=(12, 9))
    plt.scatter(X_full.iloc[:, 0], X_full.iloc[:, 2], c=X_full['anomaly_label'], cmap='coolwarm', edgecolor='k')
    plt.title("Isolation Forest")
    plt.xlabel("Wspd")
    plt.ylabel("Patv")
    plt.grid(True)

    
    # 设置刻度间隔（横坐标每 1 个单位、纵坐标每 100 个单位）
    plt.gca().xaxis.set_major_locator(MultipleLocator(1))
    plt.gca().yaxis.set_major_locator(MultipleLocator(100))

    plt.savefig(f'{filename}.png')
    plt.close()
